from skimage.feature import hog
from sklearn.svm import SVC
import numpy as np
import time, joblib
import util

train_time:float
test_time:float

class HOGSVM:
    def __init__(self):
        self.orientations = 9        # 梯度方向分箱数
        self.pixels_per_cell = (8,8) # 单元格大小
        self.cells_per_block = (2,2)
        self.svm=None

    def get_hogs(self,data:np.ndarray):
        hogs=[]
        for item in data:
            hogs.append(
                hog(item,
                    orientations=self.orientations,
                    pixels_per_cell=self.pixels_per_cell,
                    cells_per_block=self.cells_per_block,
                )
            )
        return np.vstack(hogs)

    def train(self,data:np.ndarray,label:np.ndarray):
        hogs=self.get_hogs(data)
        self.svm=SVC(kernel='linear')
        self.svm.fit(hogs,label)

    def predict(self,data:np.ndarray):
        return self.svm.predict(self.get_hogs(data))

def launch_train():
    data,label=util.load_pic("number_data/training/")
    hog_svm=HOGSVM()
    start_time=time.time()
    hog_svm.train(data,label)
    global train_time
    train_time=time.time()-start_time
    return hog_svm

def launch_test(model:HOGSVM):
    test_data,test_label=util.load_pic("number_data/testing/")
    start_time=time.time()
    predict_label=model.predict(test_data)
    global test_time
    test_time=time.time()-start_time
    util.show_result(predict_label,test_label)

def save(model:HOGSVM,path:str):
    joblib.dump(model,path)

def report_time():
    print(f"task2 trainning cost {train_time}s")
    print(f"task2 testing cost {test_time}s")

def load(path:str):
    return joblib.load(path)

def main():
    model=launch_train()
    launch_test(model)
    save(model,"result/task2/task2.joblib")
    report_time()

if __name__=="__main__":
    main()